package org.example.cases.sf1;

import com.antgroup.geaflow.api.function.base.FilterFunction;
import com.antgroup.geaflow.api.function.base.MapFunction;
import com.antgroup.geaflow.api.graph.PGraphWindow;
import com.antgroup.geaflow.api.pdata.stream.window.PWindowSource;
import com.antgroup.geaflow.api.window.impl.AllWindow;
import com.antgroup.geaflow.model.graph.edge.IEdge;
import com.antgroup.geaflow.model.graph.edge.impl.ValueEdge;
import com.antgroup.geaflow.model.graph.vertex.IVertex;
import com.antgroup.geaflow.model.graph.vertex.impl.ValueVertex;
import com.antgroup.geaflow.pipeline.Pipeline;
import com.antgroup.geaflow.pipeline.task.PipelineTask;
import com.antgroup.geaflow.view.GraphViewBuilder;
import com.antgroup.geaflow.view.IViewDesc;
import com.antgroup.geaflow.view.graph.GraphViewDesc;
import org.example.constants.Constants;
import org.example.function.CSVFileSinkV2;
import org.example.function.CSVFileSourceV2;
import org.example.model.LongStringResult;
import org.example.util.FileUtils;

import java.util.Arrays;
import java.util.Collections;

public class Case23Pipeline {

    private final String snapshotPath;

    private final Pipeline pipeline;

    public Case23Pipeline(String snapshotPath, Pipeline pipeline) {
        this.snapshotPath = snapshotPath;
        this.pipeline = pipeline;
    }

    public void submit() {
        pipeline.submit((PipelineTask) pipelineTaskCxt -> {

            CSVFileSourceV2<IVertex<Long, Double>> source = new CSVFileSourceV2<>(
                    FileUtils.pathJoin(this.snapshotPath, Constants.FILE_ACCOUNT),
                    Collections.singletonList(Constants.ACCOUNT_ID),
                    (CSVFileSourceV2.CSVLineParser<IVertex<Long, Double>>) line -> {
                        IVertex<Long, Double> vertex = new ValueVertex<>(
                                Long.valueOf(line[0]),
                                0.0
                        );
                        return Collections.singletonList(vertex);
                    }
            );

            PWindowSource<IVertex<Long, Double>> vertices = pipelineTaskCxt.buildSource(source, AllWindow.getInstance())
                    .withParallelism(Constants.SOURCE_PARALLELISM);

            PWindowSource<IEdge<Long, Double>> edges = pipelineTaskCxt.buildSource(new CSVFileSourceV2<>(
                            FileUtils.pathJoin(this.snapshotPath, Constants.FILE_ACCOUNT_TRANSFER_ACCOUNT),
                            Arrays.asList(Constants.FROM_ID, Constants.TO_ID, Constants.AMOUNT),
                            (CSVFileSourceV2.CSVLineParser<IEdge<Long, Double>>) line -> {
                                IEdge<Long, Double> edge = new ValueEdge<>(
                                        Long.valueOf(line[0]),
                                        Long.valueOf(line[1]),
                                        Double.valueOf(line[2])
                                );
                                return Collections.singletonList(edge);
                            }
                    ), AllWindow.getInstance())
                    .withParallelism(Constants.SOURCE_PARALLELISM);

            int iterationParallelism = Constants.ITERATOR_PARALLELISM;
            GraphViewDesc graphViewDesc = GraphViewBuilder
                    .createGraphView(GraphViewBuilder.DEFAULT_GRAPH)
                    .withShardNum(iterationParallelism)
                    .withBackend(IViewDesc.BackendType.Memory)
                    .build();

            PGraphWindow<Long, Double, Double> graphWindow =
                    pipelineTaskCxt.buildWindowStreamGraph(vertices, edges, graphViewDesc);

            graphWindow.compute(new Case3Algorithm.Case3VCAlgorithm(2))
                    .compute(iterationParallelism)
                    .getVertices()
                    .filter((FilterFunction<IVertex<Long, Double>>) record -> record.getValue() > 0)
                    .withParallelism(Constants.FILTER_PARALLELISM)
                    .map((MapFunction<IVertex<Long, Double>, LongStringResult>) value ->
                            new LongStringResult(value.getId(), String.format("%.2f", value.getValue())))
                    .withParallelism(Constants.MAP_PARALLELISM)
                    .sink(new CSVFileSinkV2<>("result3.csv"))
                    .withParallelism(Constants.SINK_PARALLELISM);

            graphWindow.compute(new Case2Algorithm.Case2VCAlgorithm(3))
                    .compute(iterationParallelism)
                    .getVertices()
                    .filter((FilterFunction<IVertex<Long, Double>>) record -> record.getValue() > 0)
                    .withParallelism(Constants.FILTER_PARALLELISM)
                    .map((MapFunction<IVertex<Long, Double>, LongStringResult>) value ->
                            new LongStringResult(value.getId(), "" + value.getValue().intValue()))
                    .withParallelism(Constants.MAP_PARALLELISM)
                    .sink(new CSVFileSinkV2<>("result2.csv"))
                    .withParallelism(Constants.SINK_PARALLELISM);
        });
    }
}
